use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use chrono::Utc;
use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress};
use rustls::pki_types::ServerName;
use rustls::{ClientConfig, RootCertStore};
use std::error::Error;
use std::sync::Arc;
use std::time::SystemTime;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufWriter, ReadHalf, WriteHalf};
use tokio::net::TcpStream;
use tokio_rustls::client::TlsStream;
use tokio_socks::tcp::Socks5Stream;

// 预分配缓冲区大小常量
pub const READ_BUFFER_SIZE: usize = 8192;
pub const WRITE_BUFFER_SIZE: usize = 4096;
pub const MAX_FRAME_SIZE: u64 = 10 * 1024 * 1024; // 10MB限制

pub const STACK_BUFFER_SIZE: usize = 1024;

// PerMessage-Deflate压缩器
pub struct PerMessageDeflate {
    compressor: Compress,
    decompressor: Decompress,
    server_max_window_bits: u8,
    client_max_window_bits: u8,
    compression_enabled: bool,
}

impl PerMessageDeflate {
    pub fn new(server_max_window_bits: u8, client_max_window_bits: u8) -> Self {
        let compressor =
            Compress::new_with_window_bits(Compression::default(), false, client_max_window_bits);
        let decompressor = Decompress::new_with_window_bits(false, server_max_window_bits);

        Self {
            compressor,
            decompressor,
            server_max_window_bits,
            client_max_window_bits,
            compression_enabled: false,
        }
    }

    pub fn enable_compression(&mut self, enabled: bool) {
        self.compression_enabled = enabled;
    }

    pub fn compress(&mut self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
        if !self.compression_enabled {
            return Ok(data.to_vec());
        }

        self.compressor.reset();
        let mut compressed = Vec::new();

        // 使用 compress_vec 进行压缩
        self.compressor
            .compress_vec(data, &mut compressed, FlushCompress::Sync)?;

        // 移除尾部的 0x00 0x00 0xff 0xff (per-message-deflate 规范)
        if compressed.len() >= 4 && compressed.ends_with(&[0x00, 0x00, 0xff, 0xff]) {
            compressed.truncate(compressed.len() - 4);
        }

        Ok(compressed)
    }

    pub fn decompress(&mut self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
        if !self.compression_enabled {
            return Ok(data.to_vec());
        }

        // 添加尾部标记 (per-message-deflate 规范)
        let mut data_with_tail = data.to_vec();
        data_with_tail.extend_from_slice(&[0x00, 0x00, 0xff, 0xff]);

        self.decompressor.reset(false);
        let mut decompressed = Vec::new();

        self.decompressor.decompress_vec(
            &data_with_tail,
            &mut decompressed,
            FlushDecompress::Sync,
        )?;

        Ok(decompressed)
    }
}

// 优化的握手函数 - 添加压缩扩展支持
async fn perform_handshake(
    stream: &mut TlsStream<TcpStream>,
    host: &str,
    path: &str,
    deflate: &mut PerMessageDeflate,
) -> Result<(), Box<dyn Error>> {
    let key = generate_websocket_key();

    let request = format!(
        "GET {path} HTTP/1.1\r\n\
         Host: {host}\r\n\
         Upgrade: websocket\r\n\
         Connection: Upgrade\r\n\
         Sec-WebSocket-Key: {key}\r\n\
         Sec-WebSocket-Version: 13\r\n\
         Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits\r\n\r\n",
    );

    println!("发送握手请求:\n{request}");

    // 使用BufWriter提高写入性能
    let mut buf_writer = BufWriter::with_capacity(WRITE_BUFFER_SIZE, stream);
    buf_writer.write_all(request.as_bytes()).await?;
    buf_writer.flush().await?;

    // 获取底层stream用于读取
    let stream = buf_writer.into_inner();

    // 优化响应读取 - 使用预分配缓冲区
    let mut response = Vec::with_capacity(512);
    let mut read_buf = [0u8; 256];

    loop {
        let n = stream.read(&mut read_buf).await?;
        if n == 0 {
            return Err("Connection closed during handshake".into());
        }

        response.extend_from_slice(&read_buf[..n]);

        // 检查头部结束
        if response.len() >= 4 && response.ends_with(b"\r\n\r\n") {
            break;
        }

        // 防止响应过大
        if response.len() > 4096 {
            return Err("Handshake response too large".into());
        }
    }

    let response_str = String::from_utf8_lossy(&response);
    println!("握手响应:\n{response_str}");

    if !response.starts_with(b"HTTP/1.1 101") {
        return Err(format!("握手失败: {response_str}").into());
    }

    // 检查服务器是否接受了压缩扩展
    let response_lower = response_str.to_lowercase();
    if response_lower.contains("permessage-deflate") {
        deflate.enable_compression(true);
        println!("PerMessage-Deflate 压缩已启用");
    } else {
        deflate.enable_compression(false);
        println!("服务器不支持 PerMessage-Deflate 压缩");
    }

    Ok(())
}

fn generate_websocket_key() -> String {
    let timestamp = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_nanos(); // 使用纳秒提高精度

    BASE64_STANDARD.encode(timestamp.to_string())
}

// 创建WebSocket文本帧（带掩码和可选压缩）
pub fn create_text_frame_with_compression(
    text: &str,
    deflate: &mut PerMessageDeflate,
) -> Result<Vec<u8>, Box<dyn Error>> {
    let payload = text.as_bytes();

    // 尝试压缩
    let compressed_payload = deflate.compress(payload)?;
    let is_compressed = deflate.compression_enabled && compressed_payload.len() < payload.len();

    let final_payload = if is_compressed {
        &compressed_payload
    } else {
        payload
    };
    let payload_len = final_payload.len();

    // 预分配准确大小
    let header_size = if payload_len < 126 {
        6
    } else if payload_len < 65536 {
        8
    } else {
        14
    };
    let mut frame = Vec::with_capacity(header_size + payload_len);

    // FIN=1, RSV1=(压缩标志), OPCODE=0x1 (text)
    let first_byte = if is_compressed { 0xC1 } else { 0x81 }; // RSV1=1 表示压缩
    frame.push(first_byte);

    // 长度和掩码位
    if payload_len < 126 {
        frame.push(0x80 | (payload_len as u8));
    } else if payload_len < 65536 {
        frame.push(0x80 | 126);
        frame.extend_from_slice(&(payload_len as u16).to_be_bytes());
    } else {
        frame.push(0x80 | 127);
        frame.extend_from_slice(&(payload_len as u64).to_be_bytes());
    }

    // 掩码键
    const MASK_KEY: [u8; 4] = [0x12, 0x34, 0x56, 0x78];
    frame.extend_from_slice(&MASK_KEY);

    // 掩码处理后的负载
    for (i, &byte) in final_payload.iter().enumerate() {
        frame.push(byte ^ MASK_KEY[i & 3]);
    }

    Ok(frame)
}

// 创建WebSocket文本帧（带掩码）- 保持原有接口
pub fn create_text_frame(text: &str) -> Vec<u8> {
    let payload = text.as_bytes();
    let payload_len = payload.len();

    // 预分配准确大小
    let header_size = if payload_len < 126 {
        6
    } else if payload_len < 65536 {
        8
    } else {
        14
    };
    let mut frame = Vec::with_capacity(header_size + payload_len);

    // FIN=1, OPCODE=0x1 (text)
    frame.push(0x81);

    // 长度和掩码位
    if payload_len < 126 {
        frame.push(0x80 | (payload_len as u8));
    } else if payload_len < 65536 {
        frame.push(0x80 | 126);
        frame.extend_from_slice(&(payload_len as u16).to_be_bytes());
    } else {
        frame.push(0x80 | 127);
        frame.extend_from_slice(&(payload_len as u64).to_be_bytes());
    }

    // 掩码键（保持原有逻辑）
    const MASK_KEY: [u8; 4] = [0x12, 0x34, 0x56, 0x78];
    frame.extend_from_slice(&MASK_KEY);

    // 掩码处理后的负载
    for (i, &byte) in payload.iter().enumerate() {
        frame.push(byte ^ MASK_KEY[i & 3]);
    }

    frame
}

// 解析WebSocket帧，返回(opcode, payload) - 带压缩支持
pub async fn parse_frame_with_compression<R: AsyncReadExt + Unpin>(
    reader: &mut R,
    deflate: &mut PerMessageDeflate,
) -> Result<(u8, Vec<u8>), Box<dyn Error + Send + Sync>> {
    let mut header = [0u8; 2];
    reader.read_exact(&mut header).await?;

    let opcode = header[0] & 0x0F;
    let rsv1 = (header[0] & 0x40) != 0; // 压缩标志
    let masked = (header[1] & 0x80) != 0;
    let mut payload_length = (header[1] & 0x7F) as u64;

    match payload_length {
        126 => {
            let mut len_bytes = [0u8; 2];
            reader.read_exact(&mut len_bytes).await?;
            payload_length = u16::from_be_bytes(len_bytes) as u64;
        }
        127 => {
            let mut len_bytes = [0u8; 8];
            reader.read_exact(&mut len_bytes).await?;
            payload_length = u64::from_be_bytes(len_bytes);

            if payload_length > MAX_FRAME_SIZE {
                return Err("Frame too large".into());
            }
        }
        _ => {}
    }

    if masked {
        let mut mask = [0u8; 4];
        reader.read_exact(&mut mask).await?;
    }

    // 读取payload
    let payload = if payload_length <= STACK_BUFFER_SIZE as u64 {
        let mut stack_buf = [0u8; STACK_BUFFER_SIZE];
        let len = payload_length as usize;
        reader.read_exact(&mut stack_buf[..len]).await?;
        stack_buf[..len].to_vec()
    } else {
        let mut payload = Vec::with_capacity(payload_length as usize);
        unsafe {
            payload.set_len(payload_length as usize);
        }
        reader.read_exact(&mut payload).await?;
        payload
    };

    // 如果是压缩的文本帧，进行解压缩
    let final_payload = if rsv1 && (opcode == 0x1 || opcode == 0x2) {
        match deflate.decompress(&payload) {
            Ok(decompressed) => decompressed,
            Err(e) => {
                println!("解压缩失败: {}", e);
                payload // 解压缩失败时返回原始数据
            }
        }
    } else {
        payload
    };

    Ok((opcode, final_payload))
}

// 解析WebSocket帧，返回(opcode, payload) - 保持原有接口
#[allow(clippy::uninit_vec)]
pub async fn parse_frame<R: AsyncReadExt + Unpin>(
    reader: &mut R,
) -> Result<(u8, Vec<u8>), Box<dyn Error + Send + Sync>> {
    let mut header = [0u8; 2];
    reader.read_exact(&mut header).await?;

    let opcode = header[0] & 0x0F;
    let masked = (header[1] & 0x80) != 0;
    let mut payload_length = (header[1] & 0x7F) as u64;

    match payload_length {
        126 => {
            let mut len_bytes = [0u8; 2];
            reader.read_exact(&mut len_bytes).await?;
            payload_length = u16::from_be_bytes(len_bytes) as u64;
        }
        127 => {
            let mut len_bytes = [0u8; 8];
            reader.read_exact(&mut len_bytes).await?;
            payload_length = u64::from_be_bytes(len_bytes);

            if payload_length > MAX_FRAME_SIZE {
                return Err("Frame too large".into());
            }
        }
        _ => {}
    }

    if masked {
        let mut mask = [0u8; 4];
        reader.read_exact(&mut mask).await?;
    }

    // 小payload使用栈分配，大payload使用堆分配
    let payload = if payload_length <= STACK_BUFFER_SIZE as u64 {
        let mut stack_buf = [0u8; STACK_BUFFER_SIZE];
        let len = payload_length as usize;
        reader.read_exact(&mut stack_buf[..len]).await?;
        stack_buf[..len].to_vec()
    } else {
        let mut payload = Vec::with_capacity(payload_length as usize);
        unsafe {
            payload.set_len(payload_length as usize);
        }
        reader.read_exact(&mut payload).await?;
        payload
    };

    Ok((opcode, payload))
}

/// WebSocket连接启动函数 - 带压缩支持
pub async fn start_ws_with_compression(
    domain: &'static str,
    path: Option<&'static str>,
    params: Option<&str>,
    proxy: &str,
    server_max_window_bits: u8,
    client_max_window_bits: u8,
) -> (
    ReadHalf<TlsStream<TcpStream>>,
    WriteHalf<TlsStream<TcpStream>>,
    PerMessageDeflate,
) {
    let tls_stream = if proxy.is_empty() {
        TcpStream::connect(format!("{domain}:443")).await.unwrap()
    } else {
        Socks5Stream::connect(proxy, format!("{domain}:443"))
            .await
            .unwrap()
            .into_inner()
    };

    let mut root_store = RootCertStore::empty();
    for cert in rustls_native_certs::load_native_certs().expect("could not load platform certs") {
        root_store.add(cert).unwrap();
    }

    let config = ClientConfig::builder()
        .with_root_certificates(root_store)
        .with_no_client_auth();

    let connector = tokio_rustls::TlsConnector::from(Arc::new(config));
    let current_domain = ServerName::try_from(domain).unwrap();
    let mut tls_stream = connector.connect(current_domain, tls_stream).await.unwrap();

    let mut current_params = String::new();
    if let Some(path) = path {
        current_params.push_str(path);
    }
    if let Some(params) = params {
        if path.is_some() {
            current_params.push_str(format!("?{params}").as_str());
        } else {
            current_params.push_str(format!("/?{params}").as_str());
        }
    }

    // 创建压缩器
    let mut deflate = PerMessageDeflate::new(server_max_window_bits, client_max_window_bits);

    // 握手
    perform_handshake(
        &mut tls_stream,
        domain,
        current_params.as_str(),
        &mut deflate,
    )
    .await
    .unwrap();

    println!("握手成功");

    let (reader, writer) = tokio::io::split(tls_stream);
    (reader, writer, deflate)
}

/// 原有的WebSocket连接启动函数 - 保持兼容性
pub async fn start_ws(
    domain: &'static str,
    path: Option<&'static str>,
    params: Option<&str>,
    proxy: &str,
) -> (
    ReadHalf<TlsStream<TcpStream>>,
    WriteHalf<TlsStream<TcpStream>>,
) {
    let tls_stream = if proxy.is_empty() {
        TcpStream::connect(format!("{domain}:443")).await.unwrap()
    } else {
        Socks5Stream::connect(proxy, format!("{domain}:443"))
            .await
            .unwrap()
            .into_inner()
    };

    let mut root_store = RootCertStore::empty();
    for cert in rustls_native_certs::load_native_certs().expect("could not load platform certs") {
        root_store.add(cert).unwrap();
    }

    let config = ClientConfig::builder()
        .with_root_certificates(root_store)
        .with_no_client_auth();

    let connector = tokio_rustls::TlsConnector::from(Arc::new(config));
    let current_domain = ServerName::try_from(domain).unwrap();
    let mut tls_stream = connector.connect(current_domain, tls_stream).await.unwrap();

    let mut current_params = String::new();
    if let Some(path) = path {
        current_params.push_str(path);
    }
    if let Some(params) = params {
        if path.is_some() {
            current_params.push_str(format!("?{params}").as_str());
        } else {
            current_params.push_str(format!("/?{params}").as_str());
        }
    }

    // 创建一个虚拟的压缩器用于握手（但不启用压缩）
    let mut deflate = PerMessageDeflate::new(15, 15);

    // 原有握手（不包含压缩扩展）
    let key = generate_websocket_key();
    let request = format!(
        "GET {current_params} HTTP/1.1\r\n\
         Host: {domain}\r\n\
         Upgrade: websocket\r\n\
         Connection: Upgrade\r\n\
         Sec-WebSocket-Key: {key}\r\n\
         Sec-WebSocket-Version: 13\r\n\r\n"
    );

    println!("发送握手请求:\n{request}");

    let mut buf_writer = BufWriter::with_capacity(WRITE_BUFFER_SIZE, &mut tls_stream);
    buf_writer.write_all(request.as_bytes()).await.unwrap();
    buf_writer.flush().await.unwrap();

    let mut response = Vec::with_capacity(512);
    let mut read_buf = [0u8; 256];

    loop {
        let n = tls_stream.read(&mut read_buf).await.unwrap();
        if n == 0 {
            panic!("Connection closed during handshake");
        }

        response.extend_from_slice(&read_buf[..n]);

        if response.len() >= 4 && response.ends_with(b"\r\n\r\n") {
            break;
        }

        if response.len() > 4096 {
            panic!("Handshake response too large");
        }
    }

    let response_str = String::from_utf8_lossy(&response);
    println!("握手响应:\n{response_str}");

    if !response.starts_with(b"HTTP/1.1 101") {
        panic!("握手失败: {response_str}");
    }

    println!("握手成功");

    tokio::io::split(tls_stream)
}

pub struct TimeDomain;

impl TimeDomain {
    /// 获取今天的日期
    pub fn get_date() -> String {
        Utc::now().format("%Y-%m-%d").to_string()
    }
}
